import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import eig
from typing import Tuple, List, Optional
from dataclasses import dataclass
from enum import Enum

@dataclass
class MaterialProperties:
    """Material properties for matrix and graphene."""
    E: float  # Young's modulus (GPa)
    nu: float  # Poisson's ratio
    rho: float  # Density (kg/m³)
    
class GrapheneGeometry:
    """Graphene platelet geometry properties."""
    def __init__(self, length: float = 2.5e-6, width: float = 1.5e-6, thickness: float = 1.5e-9):
        self.length = length  # Average length (m)
        self.width = width    # Average width (m) 
        self.thickness = thickness  # Average thickness (m)

class DistributionPattern(Enum):
    """Graphene distribution patterns."""
    UD = "uniform"  # Uniform distribution
    FG_V = "fg_v"   # Higher near top surface
    FG_O = "fg_o"   # Higher near top and bottom surfaces  
    FG_X = "fg_x"   # Higher in the middle

class GrapheneReinforcedComposite:
    """
    Graphene-reinforced composite material with functionally graded distribution.
    Based on the modified Halpin-Tsai model from the paper.
    """
    
    def __init__(self, matrix: MaterialProperties, graphene: MaterialProperties, 
                 geometry: GrapheneGeometry):
        self.matrix = matrix
        self.graphene = graphene
        self.geometry = geometry
        
    def volume_fraction(self, z: float, z_k: float, z_k1: float, 
                       V_G_star: float, p: float = 1.0, 
                       pattern: DistributionPattern = DistributionPattern.UD) -> float:
        """
        Calculate graphene volume fraction at position z.
        
        Args:
            z: Position through thickness
            z_k: Bottom surface coordinate of GRC layer
            z_k1: Top surface coordinate of GRC layer
            V_G_star: Total graphene volume fraction
            p: Volume fraction exponent
            pattern: Distribution pattern
            
        Returns:
            Volume fraction at position z
        """
        if pattern == DistributionPattern.UD:
            return V_G_star
        
        # Normalized coordinate
        xi = (z - z_k) / (z_k1 - z_k)
        
        if pattern == DistributionPattern.FG_V:
            return V_G_star * (xi ** p)
        elif pattern == DistributionPattern.FG_O:
            return V_G_star * (2 * xi * (1 - xi)) ** p
        elif pattern == DistributionPattern.FG_X:
            return V_G_star * (4 * xi * (1 - xi)) ** p
        else:
            return V_G_star
    
    def halpin_tsai_parameters(self, V_G: float) -> Tuple[float, float]:
        """
        Calculate Halpin-Tsai parameters.
        
        Args:
            V_G: Graphene volume fraction
            
        Returns:
            eta_L, eta_T: Longitudinal and transverse efficiency parameters
        """
        # Shape parameters
        xi_L = 2 * self.geometry.length / self.geometry.thickness
        xi_T = 2 * self.geometry.width / self.geometry.thickness
        
        # Modulus ratio
        E_ratio = self.graphene.E / self.matrix.E
        
        # Efficiency parameters
        eta_L = (E_ratio - 1) / (E_ratio + xi_L)
        eta_T = (E_ratio - 1) / (E_ratio + xi_T)
        
        return eta_L, eta_T
    
    def effective_modulus(self, V_G: float) -> float:
        """
        Calculate effective Young's modulus using modified Halpin-Tsai model.
        
        Args:
            V_G: Graphene volume fraction
            
        Returns:
            Effective Young's modulus (GPa)
        """
        if V_G == 0:
            return self.matrix.E
            
        eta_L, eta_T = self.halpin_tsai_parameters(V_G)
        
        # Longitudinal and transverse contributions
        E_L = (1 + (2 * self.geometry.length / self.geometry.thickness) * eta_L * V_G) / \
              (1 - eta_L * V_G) * self.matrix.E
        
        E_T = (1 + (2 * self.geometry.width / self.geometry.thickness) * eta_T * V_G) / \
              (1 - eta_T * V_G) * self.matrix.E
        
        # Effective modulus (3/8 longitudinal + 5/8 transverse)
        E_c = (3/8) * E_L + (5/8) * E_T
        
        return E_c
    
    def effective_properties(self, V_G: float) -> MaterialProperties:
        """
        Calculate all effective material properties.
        
        Args:
            V_G: Graphene volume fraction
            
        Returns:
            Effective material properties
        """
        # Effective Young's modulus
        E_c = self.effective_modulus(V_G)
        
        # Poisson's ratio (rule of mixtures)
        nu_c = self.graphene.nu * V_G + self.matrix.nu * (1 - V_G)
        
        # Density (rule of mixtures)
        rho_c = self.graphene.rho * V_G + self.matrix.rho * (1 - V_G)
        
        return MaterialProperties(E_c, nu_c, rho_c)
    
    def reduced_stiffness_matrix(self, V_G: float) -> np.ndarray:
        """
        Calculate reduced stiffness matrix Q for plane stress.
        
        Args:
            V_G: Graphene volume fraction
            
        Returns:
            5x5 reduced stiffness matrix
        """
        props = self.effective_properties(V_G)
        
        # Plane stress reduced stiffnesses
        Q11 = props.E / (1 - props.nu**2)
        Q22 = Q11  # Isotropic material
        Q12 = props.nu * props.E / (1 - props.nu**2)
        Q66 = props.E / (2 * (1 + props.nu))
        Q44 = Q66  # Shear modulus
        Q55 = Q66  # Shear modulus
        
        Q = np.zeros((5, 5))
        Q[0, 0] = Q11
        Q[1, 1] = Q22
        Q[0, 1] = Q[1, 0] = Q12
        Q[2, 2] = Q66
        Q[3, 3] = Q44
        Q[4, 4] = Q55
        
        return Q

# Example usage and validation
if __name__ == "__main__":
    # Material properties from Table 1 in the paper
    epoxy = MaterialProperties(E=3.0, nu=0.34, rho=1200)
    graphene = MaterialProperties(E=1010, nu=0.186, rho=1060)
    
    # Graphene geometry from Table 1
    gpl_geometry = GrapheneGeometry(
        length=2.5e-6,    # 2.5 μm
        width=1.5e-6,     # 1.5 μm  
        thickness=1.5e-9  # 1.5 nm
    )
    
    # Create GRC material
    grc = GrapheneReinforcedComposite(epoxy, graphene, gpl_geometry)
    
    # Test effective modulus calculation
    V_G_values = np.linspace(0, 0.05, 100)  # 0 to 5% volume fraction
    E_eff = [grc.effective_modulus(V_G) for V_G in V_G_values]
    
    # Plot results
    plt.figure(figsize=(10, 6))
    plt.plot(V_G_values * 100, E_eff, 'b-', linewidth=2, label='Effective Modulus')
    plt.axhline(y=epoxy.E, color='r', linestyle='--', label='Matrix Modulus')
    plt.xlabel('Graphene Volume Fraction (%)')
    plt.ylabel('Effective Young\'s Modulus (GPa)')
    plt.title('Effective Modulus vs Graphene Volume Fraction')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()
    
    print("Material property calculations implemented successfully!")
    print(f"Matrix modulus: {epoxy.E} GPa")
    print(f"Graphene modulus: {graphene.E} GPa") 
    print(f"Effective modulus at 1% GPL: {grc.effective_modulus(0.01):.2f} GPa")
    print(f"Enhancement factor: {grc.effective_modulus(0.01)/epoxy.E:.2f}")